# -*- coding: utf-8 -*-
__author__ = 'BuddyHolly'

import os
import cv2
import numpy as np
import random

ALPHA = 3
SEED = 5949


# =========================================================
# *** 图片添加水印 ***
#
# img_path_name：原图路径
# wm_path_name：水印图路径
# alpha：通道
# ran_seed：随机所使用seed
#
# 变换后的图像存在img_path_name处（更新img且后缀变为png），无返回值
# =========================================================

def en_pic(img_path_name, wm_path_name, alpha=ALPHA, ran_seed=SEED):
    # 输入原图
    img = cv2.imread(img_path_name)
    # 图像快速傅里叶变换
    img_f = np.fft.fft2(img)

    # 获得图像尺寸
    height, width, channel = np.shape(img)
    # 打开水印图
    watermark = cv2.imread(wm_path_name)
    wm_height, wm_width = watermark.shape[0], watermark.shape[1]
    x, y = range(height / 2), range(width)

    # 计算水印图是否可以添加在原图上（水印图是否小于原图）
    assert x > wm_height
    assert y > wm_width

    # 图像乱序加密（控制random seed）
    random.seed(height + width + ran_seed)
    random.shuffle(x)
    random.shuffle(y)

    # 乱序图像改变大小
    tmp = np.zeros(img.shape)
    for i in range(height / 2):
        for j in range(width):
            if x[i] < wm_height and y[j] < wm_width:
                tmp[i][j] = watermark[x[i]][y[j]]
                tmp[height - 1 - i][width - 1 - j] = tmp[i][j]
    res_f = img_f + alpha * tmp

    # 傅里叶逆变换恢复图像
    res = np.fft.ifft2(res_f)
    res = np.real(res)
    os.remove(img_path_name)
    img_path_name = img_path_name[:-4] + '.png'

    cv2.imwrite(img_path_name, res, [int(cv2.IMWRITE_JPEG_QUALITY), 100])
